﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using Kokomo.Lexical.Coding;
using Kokomo.Lexical.Expressions;
using Kokomo.PeInspector.ObjectGraphs;

namespace Kokomo.PeInspector.Reflection.Disassembly.Flow
{
	class SwitchFlowAnalyzer : FlowAnalyzer
	{
		public int ExitOffset { get { return this.BlockEndOffset; } }
		public SwitchStatement SwitchStatement { get; private set; }
		public Dictionary<int, int> CaseTable { get; private set; }
		public SwitchDefaultFlow DefaultFlow { get; private set; }

		public Expression Condition { get; private set; }
		public Type ConditionType { get; private set; }

		private static int DetermineConditionAdjustment(ref Expression condition)
		{
			BinaryOperationExpression binary = condition as BinaryOperationExpression;

			int adjustment = 0;
			if (binary != null)
			{
				if (
					(binary.Operator == ExpressionOperator.Addition)
					|| (binary.Operator == ExpressionOperator.Subtraction)
					)
				{
					IntegerConstantExpression rightConstant = binary.Right as IntegerConstantExpression;
					if (rightConstant != null)
					{
						adjustment = rightConstant.TypedValue;
						condition = binary.Left;
					}
					else
					{
						FieldReferenceExpression rightFieldRef = binary.Right as FieldReferenceExpression;
						if (
							(rightFieldRef != null)
							&& (rightFieldRef.Field.DeclaringType.IsEnum)
							)
						{
							adjustment = Convert.ToInt32(rightFieldRef.Field.GetRawConstantValue());
							condition = binary.Left;
						}
					}

					if (binary.Operator == ExpressionOperator.Subtraction)
						adjustment = -adjustment;
				}
			}

			return adjustment;
			//// Determine adjustment to switch value, if any
			//ObjectTemplate<Expression> leftTemplate = new ObjectTemplate<Expression>();
			//ObjectTemplate<ConstantExpression> rightTemplate = new ObjectTemplate<ConstantExpression>();
			//ObjectTemplate<BinaryOperationExpression> binaryTemplate = new ObjectTemplate<BinaryOperationExpression>(
			//    new MemberTemplate[] {
			//        new PropertyTemplate<Expression>(
			//            "Left",
			//            leftTemplate
			//            ),
			//        new PropertyTemplate<ConstantExpression>(
			//            "Right",
			//            rightTemplate
			//            )
			//    });

			//bool isProbablyAdjusted = binaryTemplate.Inspect(condition);
			//int switchAdjustment = 0;
			//if (isProbablyAdjusted)
			//{
			//    if (binaryTemplate.CapturedValue.Operator == ExpressionOperator.Addition)
			//    {
			//        switchAdjustment = Convert.ToInt32(rightTemplate.CapturedValue.Value);
			//        condition = leftTemplate.CapturedValue;
			//    }
			//    else if (binaryTemplate.CapturedValue.Operator == ExpressionOperator.Subtraction)
			//    {
			//        switchAdjustment = -Convert.ToInt32(rightTemplate.CapturedValue.Value);
			//        condition = leftTemplate.CapturedValue;
			//    }
			//}

			//return switchAdjustment;
		}

		public SwitchFlowAnalyzer(
			FlowAnalyzer parent,
			int blockStartOffset,
			int bodyStartOffset,
			MsilBranchTarget[] switchBranchTargets,
			MsilStackEntry conditionStackEntry
			)
			: base(parent, blockStartOffset, bodyStartOffset, bodyStartOffset, bodyStartOffset)
		{
			if (switchBranchTargets == null) throw new ArgumentNullException("switchBranchTargets");

			this.CaseTable = new Dictionary<int, int>();

			Expression condition = conditionStackEntry.Expression;
			Type conditionType = conditionStackEntry.ExpressionType;
			int switchAdjustment = DetermineConditionAdjustment(ref condition);

			SwitchStatement switchBlock = new SwitchStatement(condition);
			this.SwitchStatement = switchBlock;

			this.DefaultFlow = new SwitchDefaultFlow(
				this,
				bodyStartOffset,
				this.BodyStartOffset
				);

			this.BuildCaseTable(switchBranchTargets, conditionStackEntry);
		}

		public int? DefaultCaseStartOffset { get; internal set; }

		internal void BuildCaseTable(
			MsilBranchTarget[] switchBranchTargets,
			MsilStackEntry conditionStackEntry
			)
		{
			Expression condition = conditionStackEntry.Expression;
			Type conditionType = conditionStackEntry.ExpressionType;
			int switchAdjustment = DetermineConditionAdjustment(ref condition);

			if (this.Condition == null)
			{
				this.Condition = condition;
				this.ConditionType = conditionType;
			}
			else if (!this.Condition.Equals(condition))
			{
				throw new InvalidOperationException(Messages.SwitchConditionMismatchMessage);
			}

			for (int i = 0; i < switchBranchTargets.Length; i++)
			{
				this.CaseTable[i - switchAdjustment] = switchBranchTargets[i].TargetOffset;
			}

			int[] caseBlockStartOffsetsSorted = this.CaseTable.Values.Distinct().ToArray();
			Array.Sort(caseBlockStartOffsetsSorted);

			MsilDecompilationContext context = this.Context;

			int lastCaseOffset = caseBlockStartOffsetsSorted.Last();
			int? switchBlockEnd = null;
			for (int i = 1; i < caseBlockStartOffsetsSorted.Length; i++)
			{
				MsilBranch previousCaseExitBranch = context.GetTransitionJustBefore(caseBlockStartOffsetsSorted[i]) as MsilBranch;
				if (
					(previousCaseExitBranch != null)
					&& (previousCaseExitBranch.TargetOffset >= lastCaseOffset)
					)
				{
					switchBlockEnd =
						switchBlockEnd.HasValue
						? Math.Min(switchBlockEnd.Value, previousCaseExitBranch.TargetOffset)
						: previousCaseExitBranch.TargetOffset
						;
				}
			}
			if (!switchBlockEnd.HasValue)
			{
				switchBlockEnd = context.GetBlockEnd(lastCaseOffset);
			}

			{
				ConditionFlowAnalyzer conditionFlowParent = this.ParentAnalyzer as ConditionFlowAnalyzer;
				if (conditionFlowParent != null)
				{
					BinaryOperationExpression parentCondition = conditionFlowParent.Condition as BinaryOperationExpression;
					if (
						(parentCondition != null)
						&& (parentCondition.Left.Equals(condition))
						&& (parentCondition.Operator.Category == ExpressionOperatorCategory.Relational)
						)
					{
						conditionFlowParent.ExtendBlockTo(switchBlockEnd.Value);
						conditionFlowParent.BodyEndOffset = switchBlockEnd.Value;

						conditionFlowParent.SwitchStatement = this.SwitchStatement;
					}
				}
			}

			this.ExtendBlockTo(switchBlockEnd.Value);
			this.BodyEndOffset = switchBlockEnd.Value;
			this.BodyStartOffset = caseBlockStartOffsetsSorted[0];

			this.DefaultFlow.ExtendBlockTo(this.BodyStartOffset);
			this.DefaultFlow.BodyEndOffset = this.BodyStartOffset;
		}

		protected override void AnalyzeInnerBranch(MsilBranch branch, FlowAnalyzer innerFlow)
		{
			if (branch.TargetOffset == this.ExitOffset)
			{
				this.Context.RecordStatement(new LoopExitStatement());
				this.Context.Advance();
			}
			else
			{
				base.AnalyzeInnerBranch(branch, innerFlow);
			}
		}

		public override void Analyze()
		{
			this.DefaultFlow.Analyze();

			ILookup<int, int> caseBlockLabelLookup = this.CaseTable.ToLookup(
				entry => entry.Value,
				entry => entry.Key
				);
			int[] caseBlockStartOffsets = caseBlockLabelLookup.Select(entry => entry.Key).ToArray();
			Array.Sort(caseBlockStartOffsets);
			CaseFlowAnalyzer[] caseFlows = new CaseFlowAnalyzer[caseBlockStartOffsets.Length];
			for (int i = 0; i < caseFlows.Length - 1; i++)
			{
				int startOffset = caseBlockStartOffsets[i];
				int endOffset = caseBlockStartOffsets[i + 1];
				CaseFlowAnalyzer caseFlow = new CaseFlowAnalyzer(
					this,
					startOffset,
					endOffset
					);

				caseFlows[i] = caseFlow;

			}

			caseFlows[caseFlows.Length - 1] = new CaseFlowAnalyzer(
				this,
				caseBlockStartOffsets.Last(),
				this.BlockEndOffset
				);

			for (int i = 0; i < caseFlows.Length; i++)
			{
				CaseFlowAnalyzer caseFlow = caseFlows[i];
				IEnumerable<int> labels = caseBlockLabelLookup[caseFlow.BlockStartOffset];
				foreach (var label in labels)
				{
					Expression labelExpression = MsilDecompiler.CoerceExpression(ConstantExpression.FromValue(label), this.ConditionType, this.Context);
					caseFlow.Labels.Add(labelExpression);
				}
			}

			//ILookup<int, int> caseBlockLabelLookup = switchBranchTargets
			//    .Select(
			//        (branch, i) => new { Label = i - switchAdjustment, TargetOffset = branch.TargetOffset }
			//        )
			//    .OrderBy(branch => branch.TargetOffset)
			//    .ToLookup(
			//        (branch) => branch.TargetOffset,
			//        (branch) => branch.Label
			//        )
			//    ;

			foreach (var caseFlow in caseFlows)
			{
				this.AnalyzeInnerBlock(caseFlow);

				CaseBlock lastCaseBlock = null;
				foreach (var labelExpression in caseFlow.Labels)
				{
					lastCaseBlock = new CaseBlock(labelExpression);
					this.SwitchStatement.CaseBlocks.Add(lastCaseBlock);
				}
				if (
					this.DefaultCaseStartOffset.HasValue
					&& (caseFlow.BlockStartOffset == this.DefaultCaseStartOffset.Value)
					)
				{
					lastCaseBlock = new CaseBlock(null);
					this.SwitchStatement.CaseBlocks.Add(lastCaseBlock);
				}

				lastCaseBlock.Statements.AddRange(caseFlow.Statements);
			}
		}

		protected override void AnalyzeSwitch(MsilSwitch @switch)
		{
			base.AnalyzeSwitch(@switch);
		}
	}
}
